import numpy as np
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.tree import plot_tree

# 设置字体为支持中文的字体
plt.rcParams['font.family'] = 'SimHei'  # 或者使用 'AR PL UKai CN'
plt.rcParams['axes.unicode_minus'] = False  # 用于正常显示负号

# 生成数据
np.random.seed(np.random.randint(1, 1000))
X = np.sort(5 * np.random.rand(80, 1), axis=0)  # 生成排序后的随机特征
y = np.sin(X).ravel() + np.random.normal(0, 0.1, X.shape[0])  # 添加噪声的目标变量

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=np.random.randint(1, 1000))

# 训练决策树回归模型
tree_reg = DecisionTreeRegressor(max_depth=3,splitter="random")  # 限制树的深度以防止过拟合
tree_reg.fit(X_train, y_train)

# 预测
y_pred = tree_reg.predict(X_test)

# 评估模型
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)
print(f"测试集的均方误差 (MSE): {mse:.4f}")
print(f"测试集的 R² 分数: {r2:.4f}")

# 可视化
plt.figure(figsize=(10, 6))
# plt.scatter(X_train, y_train, color="darkorange", label="训练数据")
# plt.scatter(X_test, y_test, color="red", label="测试数据")
# plt.plot(X_test, y_pred, color="blue", linewidth=2, label="预测结果")
plot_tree(tree_reg, filled=True)
plt.xlabel("特征")
plt.ylabel("目标值")
plt.title("决策树回归")
plt.legend()
plt.show()